#ifndef __YOLO_H__
#define __YOLO_H__

#include <stdio.h>
#include <stdlib.h>
#include <sys/time.h>
#include <unistd.h>
#include <iostream>
#include <opencv2/opencv.hpp>

typedef struct{
    int label;   // 类别标签
    float score; // 类别得分
    float box_x; // 左上坐标
    float box_y; // 左上坐标
    float box_w; // 边框宽度
    float box_h; // 边框高度
}Result;

class YOLO
{
public:
    YOLO(char *model_path, int input_h, int input_w, int classes_num);
    ~YOLO();

public:
    int detect(cv::Mat &image, std::vector<Result> &results, float con_threshold, float nms_threshold);        // 进行目标检测

private:
    cv::Mat preprocess(cv::Mat &image);                                                                        // 对输入预处理
    std::vector<cv::Mat> forward(cv::Mat &input);                                                              // 进行前向传播
    std::vector<Result> postprocess(std::vector<cv::Mat> &features, float con_threshold, float nms_threshold); // 对输出后处理
    void parse_feature(cv::Mat feature, int stride, float anchor[], float con_threshold,
        std::vector<int> &labels, std::vector<cv::Rect> &bboxes, std::vector<float> &scores);                  // 解析输出特征
    float sigmoid(float x);                                                                                    // 输出激活函数

private:
    int input_h; // 模型输入高度
    int input_w; // 模型输入宽度
    int image_h; // 原始图像高度
    int image_w; // 原始图像宽度
    int rsimg_h; // 缩放图像高度
    int rsimg_w; // 缩放图像宽度
    int rspad_t; // 缩放上边填充
    int rspad_b; // 缩放下边填充
    int rspad_l; // 缩放左边填充
    int rspad_r; // 缩放右边填充

    cv::dnn::Net model;  // 网络模型结构
    int classes_num;     // 目标类别数量
    int feature_num;     // 输出特征数量
    int strides[3];      // 输出下采样率
    int batch_size;      // 输出批次数量
    int anchor_num;      // 输出锚框数量
    float anchors[3][6]; // 输出锚框尺寸
    float con_threshold; // 物体置信阈值
    float nms_threshold; // 非极大抑制值
};

#endif